from fastapi import FastAPI, UploadFile, File, HTTPException, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from typing import List, Optional, Dict, Any
import uvicorn
from dotenv import load_dotenv

from knowledge_base import knowledge_base
from rag_system import rag_system

# 加载环境变量
load_dotenv()

# 创建 FastAPI 应用
app = FastAPI(
    title="智能知识库系统",
    description="基于 FastAPI 和 RAG 的智能知识库问答系统",
    version="1.0.0"
)

# 添加 CORS 中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# 数据模型
class QueryRequest(BaseModel):
    question: str
    top_k: int = 5
    session_id: Optional[str] = None


class QueryResponse(BaseModel):
    answer: str
    source_documents: List[Dict[str, Any]]
    tokens_used: Optional[int] = None


class DocumentUploadResponse(BaseModel):
    status: str
    message: str
    chunks_count: int
    file_size: int


class KnowledgeBaseInfo(BaseModel):
    document_count: int
    collection_name: str
    vector_db_path: str


# 路由处理
@app.get("/")
async def root():
    return {"message": "智能知识库系统运行中", "status": "success"}


@app.get("/api/kb/info", response_model=KnowledgeBaseInfo)
async def get_knowledge_base_info():
    """获取知识库信息"""
    info = knowledge_base.get_collection_info()
    return info


@app.post("/api/kb/upload", response_model=DocumentUploadResponse)
async def upload_document(file: UploadFile = File(...)):
    """上传文档到知识库"""
    if not file.filename:
        raise HTTPException(status_code=400, detail="文件名不能为空")

    result = await knowledge_base.add_document(file)
    return result


@app.post("/api/kb/query", response_model=QueryResponse)
async def query_knowledge_base(request: QueryRequest):
    """查询知识库"""
    try:
        # 1. 检索相似文档
        similar_docs = knowledge_base.search_similar(
            query=request.question,
            k=request.top_k
        )

        if not similar_docs:
            return QueryResponse(
                answer="知识库中暂无相关文档，无法回答该问题。",
                source_documents=[],
                tokens_used=0
            )

        # 2. 生成答案
        result = rag_system.generate_answer(
            question=request.question,
            context_docs=similar_docs
        )

        if "error" in result:
            raise HTTPException(status_code=500, detail=result["error"])

        return QueryResponse(**result)

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"查询失败: {str(e)}")


@app.post("/api/kb/query/stream")
async def query_knowledge_base_stream(request: QueryRequest):
    """流式查询知识库"""
    from fastapi.responses import StreamingResponse
    import json

    async def generate_stream():
        try:
            # 1. 检索相似文档
            similar_docs = knowledge_base.search_similar(
                query=request.question,
                k=request.top_k
            )

            if not similar_docs:
                yield f"data: {json.dumps({'content': '知识库中暂无相关文档，无法回答该问题。', 'finished': True})}\n\n"
                return

            # 构建上下文和提示词
            context = "\n\n".join([doc.page_content for doc in similar_docs])

            system_prompt = f"""你是一个专业的知识库助手。请根据提供的上下文信息回答问题。

要求：
1. 严格基于提供的上下文信息回答
2. 如果上下文没有相关信息，请明确说明"根据现有知识库无法回答该问题"
3. 回答要准确、简洁、有用
4. 不要编造不存在的信息

上下文信息：
{context}
"""
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": request.question}
            ]

            # 2. 流式调用大模型
            stream = rag_system.client.chat.completions.create(
                model=rag_system.llm_model,
                messages=messages,
                max_tokens=2000,
                temperature=0.3,
                stream=True
            )

            # 流式返回
            for chunk in stream:
                if chunk.choices and chunk.choices[0].delta.content is not None:
                    content = chunk.choices[0].delta.content
                    yield f"data: {json.dumps({'content': content, 'finished': False})}\n\n"

            # 返回源文档信息
            sources_info = {
                'sources': [
                    {
                        'content': doc.page_content[:200] + "..." if len(doc.page_content) > 200 else doc.page_content,
                        'metadata': doc.metadata
                    } for doc in similar_docs
                ],
                'finished': True
            }
            yield f"data: {json.dumps(sources_info)}\n\n"

        except Exception as e:
            error_msg = f"流式查询失败: {str(e)}"
            yield f"data: {json.dumps({'error': error_msg, 'finished': True})}\n\n"

    return StreamingResponse(
        generate_stream(),
        media_type="text/plain",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
        }
    )


@app.delete("/api/kb/clear")
async def clear_knowledge_base():
    """清空知识库"""
    try:
        # 这里需要实现清空向量数据库的逻辑
        # 注意：这是一个危险操作，生产环境应该添加权限验证
        return {"status": "success", "message": "清空知识库功能待实现"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"清空失败: {str(e)}")


if __name__ == "__main__":
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=8000,
        reload=True
    )